import cv2
import numpy as np
import glob
import cardscan


def rectify(image, pts1):
    h, w = image.shape
    pts2 = np.float32([[0, 0], [w, 0], [0, h], [w, h]])

    m = cv2.getPerspectiveTransform(pts1, pts2)

    return cv2.warpPerspective(image, m, (w, h))

def find_corner(image):
    h, w = image.shape
    dst = cv2.cornerHarris(image,2,3,0.04)
    dst = cv2.dilate(dst, None)

    tol = 0.01*dst.max()
    min_dists = [float('inf')] * 4
    corners_image = [[0,0], [w, 0], [0, h], [w, h]]
    corners = [[0, 0] for i in range(4)]

    def distance2(pt1, pt2):
        return (pt1[0] - pt2[0])**2 + (pt1[1] - pt2[1])**2

    for j, row in enumerate(dst):
        for i, col in enumerate(row):
            if col > tol:
                pt1 = [i, j]
                for k in range(4):
                    pt2 = corners_image[k]
                    dist = distance2(pt1, pt2)
                    if dist < min_dists[k]:
                        min_dists[k] = dist
                        corners[k] = pt1
    img = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
    for corner in corners:
        cv2.circle(img, tuple(corner), 3, [255,0,0], -1)

    return corners

def prepare_image(image, agressive=False):
    corners = find_corner(image)
    rect = rectify(image, np.float32(corners))
    if agressive:
        # 比较激进的处理方法，先扩大黑色区，再收缩一点，
        # 可以去除黑块中的小白点，提高识别精度
        kernel1 = np.ones((5,5), np.uint8)
        kernel2 = np.ones((3,4), np.uint8)
        rect = cv2.erode(rect, kernel1, iterations=1)
        rect = cv2.dilate(rect, kernel2, iterations=1)
        #rect = cv2.morphologyEx(rect, cv2.MORPH_OPEN, kernel)
    return rect


if __name__ == '__main__':
    template_a = cardscan.read_template('cards/eighty-fullA.json')
    template_b = cardscan.read_template('cards/eighty-fullB.json')
    abmark = cardscan.read_abmark('cards/abmark.json')
    for filename in glob.glob('tests/*.jpg'): #['cards/tt.jpg']: #

        image = cv2.imread(filename, 0)
        new_image = prepare_image(image, True)

        cardtype=cardscan.abcheck(new_image, abmark)
        print(cardtype)
        if cardtype == "B":
            template_info = template_b
        elif cardtype == "A":
            template_info = template_a

        snum, answer = cardscan.read_sheet(new_image, template_info)
        print(snum, cardscan.convert_compact_form(answer))

        cv2.imshow('image', new_image)
        cv2.waitKey(0)
        cv2.destroyAllWindows()